package site.wanjiahao.tree;

import java.util.Comparator;


// 二叉搜索树
public class BST<E> extends BinaryTree<E> {

    private final Comparator<E> comparator;

    public BST(Comparator<E> comparator) {
        this.comparator = comparator;
    }

    public BST() {
        this(null);
    }

    public void add(E ele) {
        if (root == null) {
            // 添加root节点
            root = new Node<>(ele, null);
        } else {
            Node<E> cur = root;
            Node<E> parent = null;
            int flag = 0;
            while (cur != null) {
                int compareRes = comparator(cur.ele, ele);
                if (compareRes > 0) {
                    parent = cur;
                    cur = cur.left;
                    flag = 0;
                } else if (compareRes < 0) {
                    parent = cur;
                    cur = cur.right;
                    flag = 1;
                } else {
                    // 覆盖处理，ele可能为引用类型，更加符合实际情况
                    cur.ele = ele;
                    size++;
                    return;
                }
            }
            if (flag == 0) {
                parent.left = new Node<>(ele, parent);
            } else {
                parent.right = new Node<>(ele, parent);
            }
        }

        size++;
    }

    public int comparator(E e1, E e2) {
        if (comparator != null) {
            return comparator.compare(e1, e2);
        }
        // 直接强制转换，默认实现比较接口
        return ((Comparable<E>) e1).compareTo(e2);
    }

    public void remove(E ele) {
        // 判断当前节点的度
        Node<E> node = node(ele);

        if (hasTwoChildren(node)) {
            // 度为2：先找到对应的前驱（后继）节点，覆盖删除的节点，然后在删除对应的前驱（后继）节点（只存在度为1或者度为2）
            Node<E> predecessor = predecessor(node);
            // predecessor不存在为空
            node.ele = predecessor.ele;
            node = predecessor;
        }

        // 度为1：当前节点为父节点的什么节点，当前节点的子节点，作为这个节点存在
        Node<E> replacement = node.left != null? node.left: node.right;
        if (replacement != null) {
            replacement.parent = node.parent;
            if (node.parent != null) {
                if (node.parent.right == node) {
                    // 右节点
                    node.parent.right = replacement;
                } else {
                    // 左节点
                    node.parent.left = replacement;
                }
            } else {
                root = replacement;
            }
        }  else {
            // 度为0：直接删除
            if (node.parent != null) {
                if (node.parent.left == node) {
                    node.parent.left = null;
                } else {
                    node.parent.right = null;
                }
            } else {
                root = null;
            }
        }
        size--;
    }

    public boolean contains(E ele) {
        return node(ele) != null;
    }

    // 根据元素值，获取节点对象
    public Node<E> node(E ele) {
        Node<E> node = root;
        while (node != null) {
            int comp = comparator(node.ele, ele);
            if (comp > 0) {
                // 左边
                node = node.left;
            } else if (comp < 0) {
                // 右边
                node = node.right;
            } else {
                return node;
            }
        }
        return null;
    }

}
